#! /usr/bin/python
#coding=utf-8 
#----PYTHON FOR SIGNATRUE TESTING TOOL----
import sqlite3
import httplib
import socket
import telnetlib
import sys
import time
import re
import locale
from locale import atoi
SIG_ID = 11
EXAMPLE = 4
def usage():
	#"usage massage show"
	print("http_client.py <DIP> <DPORT> <SQLITE_DB_NAME> <TELNET_HOST_IP> [<SIGNATURE_ID>]")


def connect_sql_db(dbname):
	print("connect sqlite db name:"+dbname)
	db=sqlite3.connect(dbname)
	return db


def http_request(Dip,Dport,msg,f):
	try:
		s = socket.socket(socket.AF_INET,socket.SOCK_STREAM)
	except socket.error as elog:
		s = None
		print("creat socket error!")
		f.seek(0,2)
		f.write("--------------------------------")
		f.write("socket creating error!")
		return None
	try:
		int(Dport)
		s.settimeout(5)
		s.connect((Dip,Dport))
	except socket.error as elog:
		s.close()
		s = None
		print("socket connect error !")
		f.seek(0,2)
		f.write("--------------------------------")
		f.write("socket connecting error!")
		return None
	s.send(msg)
	print(msg)
	try:
		buf = s.recv(8192)
	except socket.error as elog:
		s.close()
		print("socket recving error!")
		f.seek(0,2)
		f.write("--------------------------------")
		f.write("socket recving error debug message:"+msg)
		s=None
		return None
	s.close()
	print (buf)
	return buf

def telnet_fortiweb(hostname):
	user ='admin'
	print("telnet "+ hostname)
	tn = telnetlib.Telnet(hostname)
	tn.read_until('login:')
	tn.write(user+'\n')
	tn.read_until('Password: ')
	tn.write('\n')
	tn.read_until('#')
	tn.write('sys sh\n')
	tn.read_until('#')
	print("Have login Fortiweb %s\n "%(hostname))
	return tn

def write_file(f,msg,flag,log):
	if(flag == True):
		m = "signature_id:" +str(msg) +"  Pass\n"
	else:
		m = "signature_id:" +str(msg) +"  Faile\n" +str(log)+"\n"
	f.seek(0,2)
	f.write("**********************\n")
	f.write(m)
	return

def change_byte_count(m,f,sig,fc):
	msg = str(m)
	tot = len(msg)
	count = 0
	col = msg.split('\r\n')
	count = len(col)
	b = 0
	if(msg.find('GET',0,3) >= 0):
		ty = msg.find('\r\n\r\n',0,len(msg) - 4)
		if(ty >= 0):
			debug = "signature:"+str(sig)+"GET header error!"+"\n"+str(msg)
			#print(debug)
			f.seek(0,2)
			f.write(debug)
			msg= msg.replace("<#","")
			msg= msg.replace("#>","")
		return msg


	for i in col:
		if(i.find('Content-Length:') >=0):
			break;

	tg  = msg.find('\r\n\r\n',0,len(msg)-4)
	if(tg > 0):
		count = len(msg) - 8 - tg
	else:
		debug = "signature:"+str(sig)+" not find header"+"\n"+str(msg)
		#print(debug)
		f.seek(0,2)
		f.write(debug)
		#return None
	
	if(msg.find('<#',tg,len(msg)-4) >= 0):
		msg= msg.replace("<#","")
		msg= msg.replace("#>","")
		count = count - 4
		if count > 0:
			new = 'Content-Length: %d'%(count)
		else :
			debug = "signature "+str(sig)+" content-lenth error:"+str(count)+"\n"+str(msg)
			fc.seek(0,2)
			fc.write(debug)
			new = i
		if(msg.find(new) < 0):
			msg = msg.replace(i,new)
			debug = "signature:"+str(sig)+" replace content-lenth "+str(i)+"->"+str(new)+"\n"
			#print(debug)
			fc.seek(0,2)
			fc.write(debug)

	else:
		msg= msg.replace("<#","")
		msg= msg.replace("#>","")
	
	print(msg)
	return msg	
	
	
def add_enter_type(m,f,sig):
	msg = m
	col = msg.split('\r\n')
	col_n = msg.split('\n')
	count = len(col)
	count_n = len(col_n)
	debug = "signature:"+str(sig)+"rn count"+str(count)+" ncount:" +str(count_n)+"\n"
	reg="^ {1,10}"
	new = ""
	for i in col_n:
		match =re.search(reg,i)
		if match:
			print("match hrere+++++"+i+'++++')
			v = re.findall(reg,i)
			t = i.replace(v[0],'',1)
			
			new = new + t + '\n'
		else :
			new = new + i + '\n'
	msg = new
		
	if(True):
		msg = msg.replace('\n','\r\n')
		msg = msg.replace('\r\r\n','\r\n')
		#print(msg)
	
	while( msg.rfind('\r\n\r\n',len(msg)-4,len(msg)) < 0):
		msg = msg + "\r\n"
	found = 0
	if(msg.find('POST',0,4) >= 0):
		if msg.find('\r\n\r\n',0,len(msg)-4) < 0:
			for j in col_n:
				if j.find('Cookie:') >= 0:
					found = 1
					break
			if found == 1:
				new = j+"\n\r"
				msg = msg.replace(j,new)

	while(msg.rfind('\r\n\r\n\r\n',len(msg)-6,len(msg)) >= 0):
		msg = msg[0:len(msg)-2]
		#msg = msg.replace('\r\n\r\n\r\n','\r\n\r\n')
	#print(msg)
	return msg


def http_packet_send(db,tn,f,dip,dport):
	count = 0
	pas = 0
	cur =db.cursor()
	if(len(sys.argv) >= 6):
		cur.execute('select * from signature_rule where chain_flag = \'0\' and \
		signature_id =\'%s\''%(str(sys.argv[5])))
	else:
		#cur.execute('select * from signature_rule where chain_flag = \'0\'')
		cur.execute('select * from signature_rule where chain_flag = \'0\' and main_class_id == \'090000000\'')
		#cur.execute('select * from signature_rule where chain_flag = \'0\' and main_class_id = \'010000000\'\
		#		or main_class_id =\'030000000\' and chain_flag = \'0\'\
		#		or main_class_id =\'050000000\' and chain_flag = \'0\'\
		#		or main_class_id =\'070000000\' and chain_flag = \'0\'\
		#		or main_class_id =\'080000000\' and chain_flag = \'0\'\
		#		or main_class_id =\'090000000\' and chain_flag = \'0\'')
		
		#cur.execute('select * from signature_rule where chain_flag = \'0\' and main_class_id = \'020000000\'\
		#		or main_class_id =\'040000000\' and chain_flag = \'0\'\
		#		or main_class_id =\'060000000\' and chain_flag = \'0\'')
	
	debug1 = open("debug1.log",'w')
	debug2 = open("debug2.log",'w')
	debug  = open("debug.log",'w')
	#print(cur.fetchone())
	end_line = str("")
	log_number = str("")
	for row in cur:
		count = count+1
		print("------------------signature id %s---------------\n"%(row[SIG_ID]))
		log =str("")
		try:
			ro = str(row[EXAMPLE])
		except UnicodeEncodeError as elog:
			ro = row[EXAMPLE].encode('utf-8')
			print("ENCODING MSG TO UTF-8" + row[SIG_ID] +"\n")

		if(len(ro) <= 0):
			print("Have no attack example,signature id is :"+row[SIG_ID]+"\n")
			pas =pas + 1
			continue
		if(ro.find('GET',0,3)>=0 or ro.find('POST',0,4)>=0):
			while(ro.find('<#')>= 0 or ro.find("#>")>= 0):
				if ro.find('<#') >=0:
					ro =ro.replace('<#','')
				if ro.find('#>') >= 0:
					ro =ro.replace('#>','')
			recvbuf = http_request(dip,dport,str(ro),debug)
		else:
			MSG = 'GET / HTTP/1.1\r\nHost:172.22.14.155\r\nAccept-Encoding:identity\r\nSigature-id:%s\r\n\r\n'%(str(row[SIG_ID]))
			recvbuf = http_request(dip,dport,MSG,debug)
		res =0

		if(recvbuf != None):
			if(recvbuf.find('HTTP/1.1 200 OK') >= 0):
				debug.seek(0,2)
				debug.write("======Signature %s REQUEST=======\n%s========RESPONSE==============%s"%(str(row[SIG_ID]),str(ro),str(recvbuf)))
			
			time.sleep(1)
			tn.write('tail -n 30 /var/log/root/alog.log \n')
			time.sleep(1)
			rmsg =str(tn.read_very_eager())
			#print(rmsg)
			if(len(sys.argv) >= 6):
				res = rmsg.find(str(row[SIG_ID]))
			else :
				if len(log_number) > 0:
					start = rmsg.find(log_number)
				else:
					print ("first package to send ,no log_number has cache")
					start = 0

				if start >= 0:
					print("----------start:%d---------------"%(start))
					res = rmsg.find(str(row[SIG_ID]),start,len(rmsg))
				else:
					print("rmsg is :%s"%(rmsg))
					res = -1;

				tn.write('tail -n 1 /var/log/root/alog.log\n')
				time.sleep(1)
				end_line = str(tn.read_very_eager())
				id_start = end_line.find("msg_id=")
				print("ID:%d,ENDLINE:%s"%(id_start,end_line))
				log_number = end_line[id_start+7:id_start+19]

			if(res < 0):
				log = "" 
			else:
				print("Find %s in rmsg"%(str(row[SIG_ID])))
		else:
			res = - 1
		

		if(res < 0):
			debug2.seek(0,2)
			debug2.write("Signature %s not found maybe Failed\n"%(str(row[SIG_ID])))
			write_file(f,row[SIG_ID],False,log)
		else:
			write_file(f,row[SIG_ID],True,log)
		#count = count+1
		#time.sleep(1)
	debug1.close()
	debug2.close()
	debug.close()
	print("pass count is :%d"%(pas))
	return count
	
def	db_check(db,f1,f2,sig):
	cur =db.cursor()
	if(sig != None):
		cur.execute('select * from signature_rule where signature_id = \'%s\''%(str(sig)))
	else:
		cur.execute('select * from signature_rule ')
	for row in cur:
		try:
			ro = str(row[EXAMPLE])
		except UnicodeEncodeError as elog:
			ro = row[EXAMPLE].encode('utf-8')
		if(ro.find("GET",0,3) < 0):
			#this package is POST or RESPONSE packet
			col = ro.split('\r\n')
			count = len(col)
			found = 0
			for i in col:
				if(i.find('Content-Length:') >=0):
					found = 1
					break;

			if(found):
				if i[15:].find("<#") >= 0:
					msg ="---------signature %s  content-length error!\n"%(row[SIG_ID])
					f2.seek(0,2)
					f2.write(msg)
					continue
				else:
					try:
						ro_count = locale.atoi(i[15:])
					except:
						msg ="-------signature %s content-lenth may have problem\n"%((row[SIG_ID]))
						f2.seek(0,2)
						f2.write(msg)
						ro_count = locale.atoi(i[16:])
				byte = ro.find('\r\n\r\n',0,len(ro)- 4) + 4
				if(byte > 0):
					count = len(ro[byte:])
					j = 0
					while ro.find('<#',byte,len(ro)) >= 0:
						j = ro.find('<#',byte,len(ro))
						count = count - 4
						byte = j+2
					new = 'Content-Length: %d'%(count)
					if ro_count != count :
						msg = "Sig %s change content-lenth %d to %d \n"%(row[SIG_ID],ro_count,count)
						f2.seek(0,2)
						f2.write(msg)
					else:
						msg = "Sig %s content lenth is correct!\n"%(row[SIG_ID])
						f1.seek(0,2)
						f1.write(msg)
						continue
				else:
					msg = "Sig %s NOt find content-lenth\n"%(row[SIG_ID])
					f2.seek(0,2)
					f2.write(msg)
					continue
			else:
				msg = "Sig %s not found content-lenth\n"%(row[SIG_ID])
				f2.seek(0,2)
				f2.write(msg)
				continue
		else:
			#packet is GET
			#print("GET MESAGE!")
			if ro.find('\r\n\r\n') >= 0:
				if len(ro.split('\r\n')) != len(ro.split('\n')):
					msg = "sig %s rn not eq n\n"%(row[SIG_ID])
					continue
				line = ro.split('\r\n')
				for i in line:
					if i.find(' ',0,1) >= 0:
						msg = "sig %s not have rnrn tailline  %s!\n"%(row[SIG_ID],i)
						f2.seek(0,2)
						f2.write(msg)

			else:
				msg = "sig %s not have rnrn tail !\n"%(row[SIG_ID])
				f2.seek(0,2)
				f2.write(msg)

	cur.close()
	return	

		



def main(dip,dport,dbname,hostname):
	#连接sqlite数据库,登陆fortiweb后台
	db = connect_sql_db(dbname)
	
	f1 = open("content_count.txt","w")
	f2 = open("packet_format.txt","w")
	if(len(sys.argv) == 6):
		db_check(db,f1,f2,sys.argv[5])
	else:
		db_check(db,f1,f2,None)
	f1.close()
	f2.close()
	print("\nTesting packet send here!\n")
	tn = telnet_fortiweb(hostname)
	#创建文件保存攻击结果记录
	f = open(Back_file,'w')
	ct = http_packet_send(db,tn,f,dip,dport)
	print("totle send packet:%d"%(ct))
	f.close()
	db.close()
	tn.write('exit\n')
	tn.read_until('#')
	tn.write('exit\n')
	tn.read_until('#')
	tn.close()
	sys.exit()





if len(sys.argv) < 5:
	usage()
	sys.exit()

Dip = str(sys.argv[1])
Dport = int(sys.argv[2])
Sql_db = str(sys.argv[3])
Back_file = str("result.txt")
Host_name = str(sys.argv[4])


if len(sys.argv) == 7:
	Bak_file = str(sys.argv[6])


if __name__ == "__main__":
	main(Dip,Dport,Sql_db,Host_name)


